package com.gitee.empty_null.spark;

import cn.hutool.core.util.StrUtil;
import com.gitee.empty_null.context.SparkChatContextHolder;
import com.gitee.empty_null.context.SparkChatDomain;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.experimental.Accessors;
import okhttp3.HttpUrl;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;

/**
 * @author xuhainan
 * @date 2023/12/12 16:04
 * @region hefei
 */
@Data
@Accessors(chain = true)
@NoArgsConstructor(staticName = "create")
public class SparkBase {

    private String hostUrl;

    private String apiKey;

    private String apiSecret;

    private static final String HEADER_FORMAT = """
            host: {}
            date: {}
            GET {} HTTP/1.1""";

    private static final String AUTHORIZATION_FORMAT = "api_key=\"{}\", algorithm=\"hmac-sha256\", headers=\"host date request-line\", signature=\"{}\"";

    private static final String paramFormat = """
            {
                "parameter": {
                    "chat": {
                        "domain": "{}",
                        "temperature": 0.5,
                        "max_tokens": 4096
                    }
                },
                "header": {
                    "app_id": "{}",
                    "uid": "{}"
                },
                "payload": {
                    "message": {
                        "text": [
                            {
                                "role": "user",
                                "content": "{}"
                            }
                        ]
                    }
                }
            }
            """;


    public String auth() throws Exception {
        URL url = new URL(hostUrl);
        String date = date();
        String requestHeader = StrUtil.format(HEADER_FORMAT, url.getHost(), date, url.getPath());
        String signature = signature(requestHeader);
        String authorization = StrUtil.format(AUTHORIZATION_FORMAT, apiKey, signature);
        // 拼接地址
        HttpUrl httpUrl = Objects.requireNonNull(HttpUrl.parse("https://" + url.getHost() + url.getPath()))
                .newBuilder()
                .addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8)))
                .addQueryParameter("date", date)
                .addQueryParameter("host", url.getHost())
                .build();
        return httpUrl.toString().replace("http://", "ws://").replace("https://", "wss://");
    }

    private static String date() {
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        return format.format(new Date());
    }


    private String signature(String param) throws Exception {
        // SHA256加密
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        byte[] hexDigits = mac.doFinal(param.getBytes(StandardCharsets.UTF_8));
        // Base64加密
        return Base64.getEncoder().encodeToString(hexDigits);
    }


    public static String getParam(String message, SparkChatDomain domain) {
        return getParam(message, domain.getType());
    }

    public static String getParam(String message, String domain) {
        String appid = SparkChatContextHolder.get().getAppid();
        String randomStr = UUID.randomUUID().toString().substring(0, 10);
        return StrUtil.format(paramFormat, domain, appid, randomStr, message);
    }
}
